#!/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
from torch_geometric.datasets import ModelNet
from torch_geometric.transforms import SamplePoints, NormalizeScale

# 定义转换：采样点并归一化
transform = [SamplePoints(1024), NormalizeScale()]

# 加载ModelNet40数据集
print("正在加载 ModelNet40 数据集...")
dataset = ModelNet(root="./data", name='40', transform=transform)

# 打印数据集信息
print(f"数据集大小: {len(dataset)}")
print(f"类别数量: {dataset.num_classes}")

# 获取第一个样本
data = dataset[0]
print("\n第一个样本信息:")
print(f"点云形状: {data.pos.shape}")  # 点云坐标
print(f"类别: {data.y.item()}")  # 类别标签
print(f"类别名称: {dataset.y_mask[data.y.item()]}")  # 类别名称

# 如果想要分开训练集和测试集
train_dataset = ModelNet(root="./data", name='40', train=True, transform=transform)
test_dataset = ModelNet(root="./data", name='40', train=False, transform=transform)

print(f"\n训练集大小: {len(train_dataset)}")
print(f"测试集大小: {len(test_dataset)}")

print("\n完成!") 